import numpy as np
from functions import *
#线性变换层
class Affine:
    def __init__(self,W,b):
        self.W=W
        self.b=b
        self.x=None
        #权重和偏置参数的导数
        self.dW=None
        self.db=None
    def forward(self,x):
        self.x=x
        out=np.dot(x,self.W)+self.b
        return out
    def backward(self,dout):
        dx=np.dot(dout,self.W.T)
        self.dW=np.dot(self.x.T,dout)
        self.db=np.sum(dout,axis=0)
        return dx


#激活函数层
class Relu:
    def __init__(self):
        self.mask=None
    def forward(self,x):
        self.mask=(x<=0)#mask 是一个数组，只有对应的 x 的元素＜0的位置才为 1，其余为 0
        out=x.copy()#一开始先让数组按照输入的来即可
        out[self.mask]=0 #把对应的元素小于0的位置数组元素设置为 0
        return out
    def backward(self,dout):
        dout[self.mask]=0 #对应的元素小于 0 的位置，导数为 0，其余位置的导数和传入的 dout 保持不变（因为是*1）
        dx=dout
        return dx

class SoftmaxWithLoss:#常用于多分类任务的输出层激活函数，这个实现的数学运算我还没明白
    def __init__(self):
        self.loss=None
        self.y=None #softmax的输出
        self.t=None #监督数据（one-hot)
    def forward(self,x,t):
        self.t=t
        self.y=softmax(x)
        self.loss=cross_entropy_error(self.y,self.t)
        return self.loss
    def backward(self,dout=1):#因为这里已经到输出的地方了，所以就是 1 了
        batch_size=self.t.shape[0]
        dx=(self.y-self.t)/batch_size #注意到除以 batch_size，传递给前面的是单个数据的误差
        return dx

class Sigmoid:
    def __init__(self):
        self.out=None
    def forward(self,x):
        self.out=sigmoid(x)
        return self.out
    def backward(self,dout):
        dx=dout*(1-self.out)*self.out
        return dx
class BatchNormalization:

    def __init__(self,gamma,beta,momentum=0.9,running_mean=None,running_var=None):
        self.gamma=gamma #缩放参数（可学习）
        self.beta=beta #偏移参数（可学习）
        self.momentum=momentum #控制统计量累积的平滑程度
        self.input_shape=None

        #测量时使用的全局统计量（训练中累计更新）
        self.running_mean=running_mean #移动平均均值
        self.running_var=running_var #移动平均方差

        #反向传播时需要用到的中间变量
        self.batch_size=None#每批次的数据量
        self.xc=None#均值
        self.std=None#标准差
        self.d_gama=None#gama的梯度
        self.d_beta=None#beta的梯度

    #主要是处理输入形状适配，让输入变成二维
    def forward(self,x,train_flg=True):
        self.input_shape=x.shape #记录输入形状,以便后期恢复
        if x.ndim!=2:#即为输入的是 4 维的情况，展平为 2 维
            N,C,H,W=x.shape
            x=x.reshape(N,-1)
        #这里才是调用核心的前向逻辑
        out=self.__forward(x,train_flg)
        return out.reshape(*self.input_shape)#再将其收回为原来的形状
    def __forward(self,x,train_flg):
        if self.running_mean is None:
            N,D=x.shape #N：批次大小，D：特征维度
            #第一次训练的时候先初始化一下均值和方差，在后面记录的时候逐渐累积计算
            self.running_mean=np.zeros(D)#移动平均均值
            self.running_var=np.zeros(D)#方差
        #代表训练模式
        if train_flg:
            #标准化
            mu=x.mean(axis=0) #按列，也就是按照每一个特征计算均值
            xc=x-mu #中心化（消除了均值偏移）
            var=np.mean(xc**2,axis=0) #按列，也就是按照特征维度计算方差
            std=np.sqrt(var+1e-7)#给 var 开根号，得到标准差，加一个极小值防止除 0
            xn=xc/std
            #存储中间变量（用于反向传播）
            self.batch_size=x.shape[0] #N批次大小
            self.xc=xc#中心化（均值为 0）
            self.xn=xn#标准化（均值为 0，方差为 1）
            self.std=std#标准差
            #更新移动平均统计量（训练中累积，测试时使用）
            self.running_mean=self.momentum*self.running_mean+(1-self.momentum)*mu
            self.running_var=self.momentum*self.running_var+(1-self.momentum)*var
        else:
            xc=x-self.running_mean
            xn=xc/np.sqrt(self.running_var+1e-7)
        #将标准化后的量进行缩放+偏移后，再像后面传递
        out=self.gamma*xn+self.beta
        return out
    #反向传播
    def backward(self,d_out):
        #如果是四维的话
        if d_out.ndim!=2:
             N,C,H,W=d_out.shape
             d_out=d_out.reshape(N,-1)
        #调用核心反向逻辑
        dx=self.__backward(d_out)
        dx=dx.reshape(*self.input_shape)
        return dx#恢复原形状的输入梯度并输出
    #这块我没推导
    def __backward(self,dout):
        #(1)计算 beta 和 gamma 的梯度
        d_beta=dout.sum(axis=0) #beta 的梯度：dout 对 beta 的偏导（求和，因为 out=gamma*n+beta）
        d_gamma=np.sum(self.xn*dout,axis=0) #gamma的偏导（x*dout求和）
        #(2)计算标准化后数据 xn 的梯度
        dxn=self.gamma*dout
        #(3)计算中心化后数据 xc 的梯度
        dxc=dxn/self.std
        #(4)计算标准差 std 的梯度（std=sqrt（var+1e-7）
        d_std=-np.sum((dxn*self.xc)/(self.std**2),axis=0)
        #(5)计算方差 var 的梯度（var=mean（xc ²） 链式法则
        d_var=0.5*d_std/self.std
        dxc +=(2.0/self.batch_size)*self.xc*d_var
        dmu=np.sum(dxc,axis=0)
        dx=dxc-dmu/self.batch_size

        self.d_gamma=d_gamma
        self.d_beta=d_beta

        return dx
